import numpy as np
import os
import torch
import torch.nn.functional as F
import math

def get_cosine_similarity(data_a, data_b):
    tensor_a = torch.tensor(data_a).view(1, -1)
    tensor_b = torch.tensor(data_b).view(1, -1)
    similarity_score = F.cosine_similarity(tensor_a, tensor_b).item()
    return  similarity_score

def handle_fin_out_data(p_data, c_data):
    p_data = p_data.reshape(-1, 9)
    c_data = c_data.reshape(-1, 9)
    c_data_sel = np.array([])
    for p in range(p_data.shape[0]):
        p_cur = p_data[p]
        score_max = 0
        c_max_id = 0
        for c in range(c_data.shape[0]):
            c_cur = c_data[c]
            if p_cur[0] == c_cur[0]:
                score_cur = get_cosine_similarity(c_cur[2:], p_cur[2:])
                if score_cur > score_max:
                    score_max = score_cur
                    c_max_id = c
        c_data_sel = np.append(c_data_sel, c_data[c_max_id])
    return c_data_sel

def check_match_result(name, sort = False, frame_id=0):
    p_path = '/home/adt/deeplearning/p_out_tensor/'
    c_path = '/home/adt/deeplearning/c_out_tensor/'

    file = name + '_' + str(frame_id) + '.bin'
    p_file_path = os.path.join(p_path, file)
    c_file_path = os.path.join(c_path, file)
    p_data = np.fromfile(p_file_path, dtype=np.float32)
    c_data = np.fromfile(c_file_path, dtype=np.float32)
    if sort:
        p_data = np.sort(p_data)
        c_data = np.sort(c_data)
    if name == 'fin_out':
        c_data = handle_fin_out_data(p_data, c_data)
        # print('p_data')
        # for i in range(int(len(p_data) / 9)):
        #     print(p_data[9 * i], p_data[9 * i + 1], p_data[9 * i + 2], p_data[9 * i + 3], p_data[9 * i + 4],p_data[9 * i + 5], p_data[9 * i + 6], p_data[9 * i + 7], p_data[9 * i + 8])
        # print('c_data')
        # for i in range(int(len(c_data) / 9)):
        #     print(c_data[9 * i], c_data[9 * i + 1], c_data[9 * i + 2], c_data[9 * i + 3], c_data[9 * i + 4],c_data[9 * i + 5], c_data[9 * i + 6], c_data[9 * i + 7], c_data[9 * i + 8])
    if len(p_data) == len(c_data):
        p_c = p_data - c_data
        score = get_cosine_similarity(p_data, c_data)
        if len(c_data) == np.sum(abs(p_c) < 1e-5):
            print(name, 'match result: success, perfectly match')
        elif score > 0.99:
            print(name, 'match result: success, score:', min(score, 1) * 100, '%')
        else:
            print(name, 'match result: failed, score:', score * 100, '%')
    else:
        print(name, 'match result: failed, p_len:', len(p_data), 'vs c_len:', len(c_data))

if __name__ == '__main__':
    check_match_result('pts')
    check_match_result('vfe_in', sort=True)
    check_match_result('vfe_out', sort=True)
    check_match_result('rpn_in')
    check_match_result('cls')
    check_match_result('reg')
    check_match_result('dir')
    check_match_result('fin_out')

